* Clear existing data
clear

* Load dataset
use "Yourpath/080224_combineddata_clean.dta", clear

* Save a backup file to prevent data loss
save "Yourpath/100824_constrained_lambda.dta", replace

* Keep only observations where phase == 2
keep if phase == 2

* Standardize variables
egen zmobility = std(mobility)
egen zlgwdp = std(lgwdp)

* Sort by fips_code and time, then generate means
sort fips_code time
by fips_code: egen mzmobility = mean(zmobility)
by fips_code: egen mzdeath = mean(zlgwdp)

* Generate deviation variables
gen dzmobility = zmobility - mzmobility
gen dzdeath = zlgwdp - mzdeath

* Generate dummy variables for cyclic
tab cyclic, gen(dcyclic)

* Generate 'used' variable
gen used = 1 if time >= 58
replace used = 0 if time < 58

* Sort by fips_code, used, and time
sort fips_code used time

* Generate lagged variables and their mean adjustments
foreach i of numlist 1/20 {
    by fips_code: gen lag`i'_zlgwdp = zlgwdp[_n - `i']
    by fips_code used: egen lag`i'_mzlgwdp = mean(lag`i'_zlgwdp)
    gen lag`i'_dzdeath = lag`i'_zlgwdp - lag`i'_mzlgwdp
}

* Generate lagged zmobility variables
foreach i of numlist 1/4 {
    by fips_code: gen lag`i'_zmobility = zmobility[_n - `i']
}

* Generate mean-adjusted zmobility
by fips_code used: egen mzmobility2 = mean(zmobility)
gen dzmobility2 = zmobility - mzmobility2

* Get a unique list of states
levelsof mst, local(states)

* Define the output Excel file path
local output_file "Yourpath/100824_state_constrained_results_lambda.xlsx"

* Initialize the Excel file and write the header
putexcel set "`output_file'", replace
putexcel A1 = "Region" B1 = "k" C1 = "Overall Impact" D1 = "Within R-squared" E1 = "Lambda Estimate" F1 = "Coefficient 1" G1 = "Coefficient 2" H1 = "Coefficient 3" I1 = "..." Y1 = "Coefficient 20"

* Define k values for the loop
local k_values 1 2 3

* Initialize row counter
local row = 2

* Loop through different k values
foreach k of local k_values {
    di "Running Constrained Model for k = `k' at National Level"
    preserve

    * Ensure no missing values in the dataset
    drop if missing(lag1_dzdeath - lag20_dzdeath) | missing(zmobility)

    * Use the given initial parameters
    local b0_initial = -0.5
    local b1_initial = 0.3

    * Construct lag coefficient expression
    local expr "zmobility = "
    forval n = 1/20 {
        local t = `n'
        if `k' == 1 {
            local beta_t = "(-exp({b0}))*exp({lambda})*exp(-exp({lambda})*(`t' - 1))"
        }
        else if `k' == 2 {
            local beta_t = "(-exp({b0}))*exp({lambda})^2*`t'*exp(-exp({lambda})*(`t' - 1))"
        }
        else if `k' == 3 {
            local beta_t = "(-exp({b0}))*exp({lambda})^3*(`t'^2)*exp(-exp({lambda})*(`t' - 1))/2"
        }
        local expr "`expr' `beta_t' * lag`n'_dzdeath + "
    }
    local expr = substr("`expr'", 1, length("`expr'") - 3)

    * Display expression for debugging
    di "`expr'"

    * Run the constrained nl model with error handling
    capture noisily nl (`expr'), initial(b0 `b0_initial' lambda `b1_initial') iterate(200)
    if _rc != 0 {
        di as error "Model did not converge for National level at k = `k'. Skipping..."
        restore
        continue
    }

    * Get parameter estimates
    scalar b0_est = -exp(_b[/b0])
    scalar lambda_est = exp(_b[/lambda])

    * Calculate coefficients and overall impact
    scalar overall_impact_constrained = 0
    forval n = 1/20 {
        local t = `n'
        if `k' == 1 {
            local coef`n' = b0_est * lambda_est * exp(-lambda_est * (`t' - 1))
        }
        else if `k' == 2 {
            local coef`n' = b0_est * (lambda_est^2) * `t' * exp(-lambda_est * (`t' - 1))
        }
        else if `k' == 3 {
            local coef`n' = b0_est * (lambda_est^3) * (`t'^2) * exp(-lambda_est * (`t' - 1)) / 2
        }
        scalar overall_impact_constrained = overall_impact_constrained + `coef`n''
    }
    di "Overall Impact (Constrained - National, k=`k'): " overall_impact_constrained

    * Generate predicted values
    capture drop pred_mob_d_constrained
    gen pred_mob_d_constrained = 0
    forval n = 1/20 {
        replace pred_mob_d_constrained = pred_mob_d_constrained + `coef`n'' * lag`n'_dzdeath
    }

    * Run fixed-effects model to get R-squared
    xtreg zmobility pred_mob_d_constrained spring fall summer i.cyclic, fe
    scalar r_squared_constrained = e(r2_w)
    di "Within R-squared (National Constrained, k=`k'): " r_squared_constrained

    * Write results to the Excel file
    putexcel A`row' = "National" B`row' = `k' C`row' = overall_impact_constrained D`row' = r_squared_constrained E`row' = lambda_est
    forval n = 1/20 {
        putexcel `=char(70 + `n' - 1)'`row' = `coef`n''
    }

    * Increment row counter
    local row = `row' + 1

    restore

    * Run constrained model for each state
    foreach state of local states {
        di "Running Constrained Model for state: `state', k = `k'"
        preserve
        keep if mst == "`state'"

        * Ensure no missing values in the dataset
        drop if missing(lag1_dzdeath - lag20_dzdeath) | missing(zmobility)

        * Run the constrained nl model with error handling
        capture noisily nl (`expr'), initial(b0 `b0_initial' lambda `b1_initial') iterate(200)
        if _rc != 0 {
            di as error "Model did not converge for state `state' at k = `k'. Skipping..."
            restore
            continue
        }

        * Get parameter estimates
        scalar b0_est = -exp(_b[/b0])
        scalar lambda_est = exp(_b[/lambda])

        * Calculate coefficients and overall impact
        scalar overall_impact_constrained_state = 0
        forval n = 1/20 {
            local t = `n'
            if `k' == 1 {
                local coef`n' = b0_est * lambda_est * exp(-lambda_est * (`t' - 1))
            }
            else if `k' == 2 {
                local coef`n' = b0_est * (lambda_est^2) * `t' * exp(-lambda_est * (`t' - 1))
            }
            else if `k' == 3 {
                local coef`n' = b0_est * (lambda_est^3) * (`t'^2) * exp(-lambda_est * (`t' - 1)) / 2
            }
            scalar overall_impact_constrained_state = overall_impact_constrained_state + `coef`n''
        }
        di "Overall Impact (Constrained - State: `state', k=`k'): " overall_impact_constrained_state

        * Generate predicted values
        capture drop pred_mob_d_constrained
        gen pred_mob_d_constrained = 0
        forval n = 1/20 {
            replace pred_mob_d_constrained = pred_mob_d_constrained + `coef`n'' * lag`n'_dzdeath
        }

        * Run fixed-effects model to get R-squared
        xtreg zmobility pred_mob_d_constrained spring fall summer i.cyclic, fe
        scalar r_squared_constrained_state = e(r2_w)
        di "Within R-squared (State: `state', k=`k'): " r_squared_constrained_state

        * Write state-level results to the Excel file
        putexcel A`row' = "`state'" B`row' = `k' C`row' = overall_impact_constrained_state D`row' = r_squared_constrained_state E`row' = lambda_est
        forval n = 1/20 {
            putexcel `=char(70 + `n' - 1)'`row' = `coef`n''
        }

        * Increment row counter
        local row = `row' + 1

        restore

